import os
import argparse

from statsmodels.stats.multitest import multipletests

from utils import gopen,TitleParser

def get_p_value(filename,pvalue_col):
    """
    get pvalue from a file
    """
    pvalue_list = []
    lines = []
    with gopen(filename,"r") as indata:
        title = indata.readline()
        tp = TitleParser(title)
        for line in indata:
            line_list = line.strip().split('\t')
            pvalue = tp.get_field(line_list,pvalue_col)
            pvalue_list.append(float(pvalue))
            lines.append(line_list)
    return pvalue_list,lines,title

def calibrate_p_value(pvalue_list,methods):
    """
    adjustment of pvalue by method
    """
    methods = methods.split(',')
    adjust_p_list = []
    for m in methods:
        adjust_p = multipletests(pvalue_list,alpha=0.05,method=m)
        adjust_p_list.append(adjust_p[1])
    return adjust_p_list

def output_p_value(adjust_p_list,lines,title,methods,ofile):
    """
    output adjusted p value 
    """
    otitle = title.strip() + "\t" + methods.replace(',',"\t") + '\n'
    with gopen(ofile,'w') as odata:
        odata.write(otitle)
        for line,adjust_p in zip(lines,zip(*adjust_p_list)):
            oline = line + [str(i) for i in adjust_p]
            odata.write("\t".join(oline)+"\n")
    print("adjustment of p value output to {}".format(ofile))

def main(args):
    """
    multiple test calibration 
    """
    infile = args.get('infile')
    pvalue_col = args.get('pvalue_column')
    methods = args.get("method")
    ofile = args.get('ofile')
    pvalue_list,lines,title = get_p_value(infile,pvalue_col)
    adjust_p_list = calibrate_p_value(pvalue_list,methods)
    output_p_value(adjust_p_list,lines,title,methods,ofile)

if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
        contact: chenming@novogene.com
        """,
        description="""
    Methods that can be used to execute multiple test calibration:
    bonferroni : one-step correction
    sidak : one-step correction
    holm-sidak : step down method using Sidak adjustments
    holm : step-down method using Bonferroni adjustments
    simes-hochberg : step-up method (independent)
    hommel : closed method based on Simes tests (non-negative)
    fdr_bh : Benjamini/Hochberg (non-negative)
    fdr_by : Benjamini/Yekutieli (negative)
    fdr_tsbh : two stage fdr correction (non-negative)
    fdr_tsbky : two stage fdr correction (non-negative)
    """)
    parser.add_argument("--infile",'-i',help="file required to add multiple test calibration with first line is title name")
    parser.add_argument("--pvalue_column",'-p',help="which column is p value required to calibration")
    parser.add_argument("--method",'-m',default="bonferroni,fdr_bh,fdr_by",help="method to use calibration,default is bonferroni,fdr_bh,fdr_by to see more selection by program help")
    parser.add_argument("--ofile",'-o',help="ofile name")
    args = parser.parse_args()
    main(vars(args))